"""
基于CSV文件的数据源实现
"""
import os
import pandas as pd
from .pandas_feed import PandasFeed


class CSVFeed(PandasFeed):
    """
    基于CSV文件的数据源
    
    从CSV文件读取数据，继承自PandasFeed
    """
    
    def __init__(self, datetime_column='datetime', datetime_format=None):
        """
        初始化CSVFeed
        
        Parameters
        ----------
        datetime_column : str, optional
            日期时间列的名称，默认为'datetime'
        datetime_format : str, optional
            日期时间的格式字符串，如果为None则自动推断
        """
        super().__init__(datetime_column=datetime_column)
        self._datetime_format = datetime_format
        
    def add_csv_file(self, instrument, filepath, datetime_column=None, datetime_format=None, **kwargs):
        """
        添加CSV文件数据
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        filepath : str
            CSV文件路径
        datetime_column : str, optional
            日期时间列的名称，如果为None则使用初始化时指定的列名
        datetime_format : str, optional
            日期时间的格式字符串，如果为None则使用初始化时指定的格式
        **kwargs
            传递给pandas.read_csv的额外参数
        """
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"找不到文件: {filepath}")
            
        if datetime_column is not None:
            self._datetime_column = datetime_column
            
        dt_format = datetime_format if datetime_format is not None else self._datetime_format
        
        # 读取CSV文件
        df = pd.read_csv(filepath, **kwargs)
        
        # 转换日期时间列
        if self._datetime_column in df.columns:
            df[self._datetime_column] = pd.to_datetime(df[self._datetime_column], format=dt_format)
            
        # 添加到数据源
        self.add_data_frame(instrument, df, self._datetime_column)
        
    def add_csv_dir(self, directory, filename_pattern=None, instrument_from_filename=True, 
                   datetime_column=None, datetime_format=None, **kwargs):
        """
        从目录中添加多个CSV文件
        
        Parameters
        ----------
        directory : str
            包含CSV文件的目录路径
        filename_pattern : str, optional
            文件名匹配模式，如果为None则读取所有.csv文件
        instrument_from_filename : bool, optional
            是否从文件名提取品种代码，默认为True
        datetime_column : str, optional
            日期时间列的名称，如果为None则使用初始化时指定的列名
        datetime_format : str, optional
            日期时间的格式字符串，如果为None则使用初始化时指定的格式
        **kwargs
            传递给pandas.read_csv的额外参数
        """
        if not os.path.isdir(directory):
            raise NotADirectoryError(f"找不到目录: {directory}")
            
        # 遍历目录中的文件
        for filename in os.listdir(directory):
            if filename.endswith('.csv'):
                if filename_pattern is None or filename_pattern in filename:
                    filepath = os.path.join(directory, filename)
                    
                    # 从文件名提取品种代码
                    if instrument_from_filename:
                        instrument = os.path.splitext(filename)[0]
                    else:
                        # 如果不从文件名提取，则需要在CSV中有一个instrument列
                        instrument = None
                        
                    self.add_csv_file(instrument, filepath, datetime_column, datetime_format, **kwargs)
